flax
jax
matplotlib
tensorflow
numpy
absl-py
optax
